import os
import sys
import base64


def doc_to_base64(input_file, output_file=None):
    """
    将Word文档转换为Base64编码并保存到文件
    
    Args:
        input_file (str): 输入的Word文档路径
        output_file (str): 输出的Base64文件路径，如果为None则自动生成
    """
    try:
        # 检查输入文件是否存在
        if not os.path.exists(input_file):
            print(f"错误：文件 '{input_file}' 不存在")
            return False
        
        # 读取Word文档的二进制数据
        with open(input_file, 'rb') as file:
            file_data = file.read()
        
        # 将二进制数据编码为Base64
        base64_encoded = base64.b64encode(file_data).decode('utf-8')
        
        # 生成输出文件名（如果未提供）
        if output_file is None:
            file_name = os.path.splitext(os.path.basename(input_file))[0]
            output_file = f"{file_name}_base64.txt"
        
        # 将Base64数据保存到文件
        with open(output_file, 'w', encoding='utf-8') as file:
            file.write(base64_encoded)
        
        print(f"转换成功！")
        print(f"输入文件: {input_file}")
        print(f"输出文件: {output_file}")
        print(f"Base64数据长度: {len(base64_encoded)} 字符")
        
        return True
        
    except Exception as e:
        print(f"转换过程中发生错误: {str(e)}")
        return False
        
        
def base64_to_doc(input_file, output_file=None):
    """
    将Base64编码文件转换回Word文档
    
    Args:
        input_file (str): 输入的Base64文件路径
        output_file (str): 输出的Word文档路径，如果为None则自动生成
    """
    try:
        # 检查输入文件是否存在
        if not os.path.exists(input_file):
            print(f"错误：文件 '{input_file}' 不存在")
            return False
        
        # 读取Base64数据
        with open(input_file, 'r', encoding='utf-8') as file:
            base64_data = file.read()
        
        # 去除可能的空白字符
        base64_data = base64_data.strip()
        
        # 将Base64数据解码为二进制
        file_binary = base64.b64decode(base64_data)
        
        # 生成输出文件名（如果未提供）
        if output_file is None:
            file_name = os.path.splitext(os.path.basename(input_file))[0]
            # 尝试从Base64数据中推断文件类型
            if file_binary[:4] == b'PK\x03\x04':  # ZIP文件签名（docx格式）
                output_file = f"{file_name}.docx"
            elif file_binary[:8] == b'\xD0\xCF\x11\xE0\xA1\xB1\x1A\xE1':  # OLE文件签名（doc格式）
                output_file = f"{file_name}.doc"
            else:
                output_file = f"{file_name}_restored.docx"  # 默认使用docx
        
        # 将二进制数据保存为Word文档
        with open(output_file, 'wb') as file:
            file.write(file_binary)
        
        print(f"转换成功！")
        print(f"输入文件: {input_file}")
        print(f"输出文件: {output_file}")
        print(f"文件大小: {len(file_binary)} 字节")
        
        return True
        
    except Exception as e:
        print(f"转换过程中发生错误: {str(e)}")
        return False


def usage():
    print("用法: python doc_converter.py <doc2base|base2doc> <输入文件> [输出文件]")
    print("示例: python doc_converter.py doc2base xxx.docx")
    print("示例: python base64_to_doc.py base2doc xxx.txt")
    
    
if __name__ == "__main__":
    """主函数"""
    if len(sys.argv) < 3:
        usage()
        exit()

    func = sys.argv[1]
    input_file = sys.argv[2]
    output_file = sys.argv[3] if len(sys.argv) > 3 else None
    if func not in ['doc2base', 'base2doc']:
        usage()
        exit()
    
    if func == 'doc2base':
        doc_to_base64(input_file, output_file)
    elif func == 'base2doc':
        base64_to_doc(input_file, output_file)
    else:
        usage()
        exit()
    